import collections
import re
import cPickle
import numpy
import os
import base64

try:
    from __init__ import merge_feature_inplace
except:
    # run with test
    import sys
    sys.path.append(os.path.dirname(__file__))

def gff_line(gline):
    """takes the text line and returns a feature dictionary"""
    if gline[0] == '#': return None
    line = gline.rstrip().split('\t')
    ftype = line[2]#.upper()
    feat = dict(
            start = int(line[3])
            , stop  = int(line[4])
            , type  = ftype
            , strand= line[6] in ('-', '-1') and -1 or 1
            , chr   = line[0]
            , locations = {}
            , attrs = {}
            )
    feat['locations'][ftype] = [(feat['start'], feat['stop'])]

    attrs = {}
    if "=" in line[8]:
        for pair in line[8].split(';'):
            if not pair: continue
            pair = pair.split('=')
            attrs[pair[0]] = pair[1]
    if attrs == {}: attrs["ID"] = line[8]
    # added rname to keep the attr that we actually want to use. 
    feat['name'] = attrs.get('rname', attrs.get("Parent", attrs.get("ID", attrs.get("Name"))))
    feat['attrs'] = attrs
    return feat

def gff_feat_to_str(f):
    attrs = ";".join(["%s=%s" % (k,v) for k, v in f['attrs'].iteritems()])
    strand = f['strand'] in (1, '1', '+') and '+' or '-'
    return "\t".join(map(str, (f['chr'], 'ucb', f.get('type', 'gene'), f['start'], f['stop'], ".", strand, ".", attrs)))


def read_gff(gff_file, chrn=None, pickle=True, merge=False):
    """
    a function that takes a gff file, parses it, and returns a Genedex-able object of the
    requested chromosome.
    gff_file: ...

    chrn: and identifier for chromosome that matches the first column in gff_file, if None
          return all data.

    pickle: if true, the parsing will occur only once, and the result
            will be saved in a pickle to allow for quick extraction of other
            chromosomes.

    merge: combine overlapping exons -- e.g. when there are
    alternative splicings. see docstring for genedex.misc.merge_feature_inplace()

    """
    if pickle:
        serialized_args = base64.urlsafe_b64encode(gff_file + str(int(merge)))
        extension = "." + serialized_args + ".gkl"
    # there's an up to date pickle, just return the requested chr.
    if pickle and os.path.exists(gff_file + extension) \
                  and os.stat(gff_file + extension).st_mtime \
                  >=  os.stat(gff_file            ).st_mtime:
        data = cPickle.load(open(gff_file + extension, 'rb'))
        if not chrn: return data
        return data[chrn]
        

    gff_by_chr = collections.defaultdict(dict)
    by_name = {}

    for line in open(gff_file):

        feat = gff_line(line)
        if not feat: continue

        # so in a gff file, each exon/cds, etc is stored in a separate
        # line. here, it merges these by feature name -- 
        # appending the locs to the appropriate feature type.
        if feat['name'] in by_name:
            existing = by_name[feat['name']]
            if feat['start'] < existing['_min']: existing['_min'] = feat['start']
            if feat['stop']  > existing['_max']: existing['_max'] = feat['stop']
            if feat['type'] in existing['locations']:
                existing['locations'][feat['type']].extend(feat['locations'][feat['type']])
            else:
                existing['locations'][feat['type']] = feat['locations'][feat['type']]
        else:
            feat['_min'] = feat['start']
            feat['_max'] = feat['stop']
            by_name[feat['name']] = feat

    # find the overall start, stop
    for feature in by_name.itervalues():
        feature['start'] = feature.pop('_min')
        feature['stop']  = feature.pop('_max')
        del feature['type']
        if merge: # this also sorts.
            merge_feature_inplace(feature)
        else:
            for v in feature['locations'].values(): v.sort()

        gff_by_chr[feature['chr']][feature['name']] = feature

    if pickle:
        cPickle.dump(gff_by_chr, open(gff_file + extension, "wb"), -1)
    if not chrn: return gff_by_chr
    return gff_by_chr[chrn]


if __name__ == "__main__":
    import doctest
    doctest.testmod()
